#!/home/agilex/miniconda3/envs/graspnet/bin/python
# -- coding: UTF-8
import rospy
from tf2_ros import TransformBroadcaster
from geometry_msgs.msg import TransformStamped
import numpy as np

import struct
import std_msgs
from sensor_msgs.msg import PointCloud2, PointField

def rad2deg(rad):
    return rad / np.pi * 180. 

def open3d2msg(open3d_cloud):
    points = np.asarray(open3d_cloud.points, dtype=np.float32)
    colors = (np.asarray(open3d_cloud.colors) * 255).astype(np.uint8) if len(open3d_cloud.colors) > 0 else np.full((points.shape[0], 3), 255, dtype=np.uint8)

    # Prepare data for PointCloud2
    cloud_data = np.zeros((points.shape[0], 4), dtype=np.float32)
    cloud_data[:, 0:3] = points

    for i in range(colors.shape[0]):
        r, g, b = colors[i]
        rgb_packed = (r << 16) | (g << 8) | b
        rgb_float = struct.unpack('f', struct.pack('I', rgb_packed))[0]
        cloud_data[i, 3] = rgb_float

    # Flatten data to byte array
    cloud_data = cloud_data.flatten().tobytes()

    # Create PointCloud2 message
    header = std_msgs.msg.Header()
    header.stamp = rospy.Time.now()
    header.frame_id = 'cam_link'

    fields = [
        PointField('x', 0, PointField.FLOAT32, 1),
        PointField('y', 4, PointField.FLOAT32, 1),
        PointField('z', 8, PointField.FLOAT32, 1),
        PointField('rgb', 12, PointField.FLOAT32, 1),
    ]

    point_cloud_msg = PointCloud2(
        header=header,
        height=1,
        width=points.shape[0],
        is_dense=False,
        is_bigendian=False,
        fields=fields,
        point_step=16,
        row_step=16 * points.shape[0],
        data=cloud_data
    )
    return point_cloud_msg

class PublishNode:
    def __init__(self) -> None:
        rospy.init_node("grasp_node")
        self.tf_broad = TransformBroadcaster()
        self.msg_tf = TransformStamped()
        self.rate = rospy.Rate(10)

        self.pub_cloud = rospy.Publisher("/points", PointCloud2, queue_size=10)
    
    def publish_tf(self, frame_id, child_frame_id, q, t):
        self.msg_tf.header.frame_id = frame_id
        self.msg_tf.child_frame_id = child_frame_id
        # -0.002 -0.015  0.201 -0.295 -0.293 -0.645  0.641
        self.msg_tf.transform.rotation.x = q[0]
        self.msg_tf.transform.rotation.y = q[1]
        self.msg_tf.transform.rotation.z = q[2]
        self.msg_tf.transform.rotation.w = q[3]
        self.msg_tf.transform.translation.x = t[0]
        self.msg_tf.transform.translation.y = t[1]
        self.msg_tf.transform.translation.z = t[2]
       
        self.msg_tf.header.stamp = rospy.Time.now()
        tf_broad.sendTransform(msg_tf)
            
    def publish_cloud(self, cloud):
        msg_cloud = open3d2msg(cloud)
        self.pub_cloud.publish(msg_cloud)


if __name__ == "__main__":
    rospy.init_node("tf_node")
    tf_broad = TransformBroadcaster()
    msg_tf = TransformStamped()

    msg_tf.header.frame_id = "cam_link"
    msg_tf.child_frame_id = "object_link"
    # -0.002 -0.015  0.201 -0.295 -0.293 -0.645  0.641
    msg_tf.transform.rotation.x = -0.295
    msg_tf.transform.rotation.y = -0.293
    msg_tf.transform.rotation.z = -0.645
    msg_tf.transform.rotation.w = 0.641
    msg_tf.transform.translation.x = -0.002
    msg_tf.transform.translation.y = -0.015 
    msg_tf.transform.translation.z = 0.201
    rate = rospy.Rate(20)
    
    while not rospy.is_shutdown():
        msg_tf.header.stamp = rospy.Time.now()
        tf_broad.sendTransform(msg_tf)
        rate.sleep()